# environment.py
from code1.uav import UAV
import numpy as np
import math
import ipdb
energy_consume = 1

class GroundUser:
    def __init__(self, position):
        self.position = position  # 地面用户位置 (x, y)
        self.communication_demand = np.random.randint(10, 50)  # 随机生成通信需求
        self.served = False

    def __repr__(self):
        return f"Gu(position={self.position}, demand={self.communication_demand})"

class MultiUAVEnv:
    def __init__(self, num_uavs, energy_states, cruise_states, grid_size, num_gu,num_actions=4, charging_stations=[(500, 500),(200,200),(800,800)]):
        # 初始化多无人机环境
        self.uavs = [
            UAV(energy_states, cruise_states, num_actions,charging_stations, id = i, init_energy=100, init_position=(i*200+200, i*300+300), grid_size=grid_size)
            for i in range(num_uavs)
        ]
        self.grid_size = grid_size  # 网格大小，假设是正方形
        self.energy_states = energy_states  # 用于限制能量状态范围
        self.cruise_states = cruise_states  # 用于限制巡航状态范围
        self.target_area = grid_size ** 2  # 目标区域的总网格数
        self.charging_stations = charging_stations  # 充电站位置
        
        # 初始化地面用户列表
        self.ground_users = [GroundUser((np.random.randint(0, grid_size), np.random.randint(0, grid_size))) for _ in range(num_gu)]
        

        # 打印出所有地面用户的位置
        print("Ground Users:", self.ground_users)
        print("Charging Stations:", self.charging_stations)

    def get_state(self, uav):
        # 限制能量状态的值在范围内
        energy_state = min(max(int(uav.energy / 10), 0), self.energy_states - 1)
        
        # 将位置 (x, y) 转换为巡航状态 (cruise_state)
        x, y = uav.position
        cruise_state = min(max(y * self.grid_size + x, 0), self.cruise_states - 1)
        
        return energy_state, cruise_state
    def perform_action(self, uav, energy_action, cruise_action):

        reward = 0
        reward_energy = 0
        charging_accept = True
        
        charging_uavs =  [other_uav for other_uav in self.uavs if other_uav.position in self.charging_stations and other_uav != uav]
        
        if len(charging_uavs) == len(self.uavs) -1:
            charging_accept = False
        else:
            charging_accept = True
            
        # 检查无人机是否在充电站位置，且选择了充电动作
        if uav.position in self.charging_stations and energy_action and charging_accept== 0:
            uav.energy = 100  # 重置电量为满电状态
            print(f"UAV{uav.id} has already recharged and cannot perform other actions.")
            next_energy_state, next_cruise_state = self.get_state(uav)
            return -1,0, next_energy_state, next_cruise_state # 返回0奖励和当前状态
        else:
            uav.reduce_energy(energy_consume)  # 每个非充电动作都消耗固定电量
            reward -= energy_consume
            
        if energy_action:
            print(f"UAV {uav.id} at {uav.position} choose the energy action 1. Heading to nearest charging station.")
            uav.go_to_nearest_charging_station ()
        else:
            # 执行飞行控制动作
            uav.execute_cruise_action(cruise_action)
            reward_energy += 1
                
        service_providing_nows = []
        
        
        # 检查是否有用户在当前位置，并提供服务
        service_reward = 0
        for gu in self.ground_users:
            if not gu.served and np.linalg.norm(np.array(gu.position)-np.array(uav.position))<= 100:
                print(f"UAV {uav.id} at {uav.position} is near user at {gu.position} with demand {gu.communication_demand}")
                service_reward = uav.serve_ground_user(gu.position,gu.communication_demand)
                if service_reward == 1:
                    #reward = reward + gu.communication_demand/2
                    gu.served = True
                    service_providing_nows.append(gu.communication_demand)
                    print(f"UAV {uav.id} at {uav.position} successfully served user at {gu.position}")
                else:
                    service_providing_nows.append(0)
            else:
                    service_providing_nows.append(0)
        
        if len(service_providing_nows) == 0:
            for gu in self.ground_users:
                service_providing_nows.append(0)
        
        for other_uav in self.uavs:
            if other_uav != uav:
                if np.linalg.norm(np.array(other_uav.position)-np.array(uav.position))<= 100:
                    print(f"UAV {uav.id} at {uav.position} is near UAV {other_uav.id} at {other_uav.position}")
                    reward -= 5
        
        # 获取执行动作后的新状态
        next_energy_state, next_cruise_state = self.get_state(uav)
        
        fairness_reward = self.calculate_fairness_reward(service_providing_nows)
        other_reward = self.calculate_other_reward(service_providing_nows)
        reward += fairness_reward*other_reward/energy_consume
        #print(f"the fainess reward is {fairness_reward} and the other reward is {other_reward}")
        
        
        return reward_energy,reward, next_energy_state, next_cruise_state
    
    
    def calculate_other_reward(self,service_providing_nows):
        other_reaward = 0
        fenmu = 0
        for uav_served in service_providing_nows:
            if uav_served > 0:
                fenmu += 1
        if fenmu == 0:
            return 0
        for uav_served in service_providing_nows:
            fenzi = uav_served
            other_reaward += fenzi/fenmu
            
        return other_reaward
            
            
    
    
    def calculate_fairness_reward(self,service_providing_nows):
        # 计算服务公平性奖励
        sum_demand = 0
        b_k = []
        mu_k = []
        b_k_sum = 0
        beta_k_sum = 0
        beta_k = []
        fenzi = 0
        fenmu = 0
        i = 0
        for gu in self.ground_users:
            sum_demand += gu.communication_demand
            
        for gu in self.ground_users:
            b_k.append(gu.communication_demand/sum_demand)
            b_k_sum +=gu.communication_demand/sum_demand
        
        for gu,bk in zip(self.ground_users,b_k):
            mu_k1 = (1+1/(math.log(b_k[1]/b_k_sum)))/(1+1/(math.log(bk/b_k_sum)))
            mu_k.append(mu_k1)
        
        
        for gu_served in service_providing_nows:
            beta_k_sum += gu_served
            
            
        if beta_k_sum >0:
            print("kaishitiaoshi")
            
        for gu_served in service_providing_nows:
            if beta_k_sum == 0:
                beta_k.append(0)
            else:
                beta_k.append(gu_served/beta_k_sum)
            
        for gu in self.ground_users :
            if beta_k[i] == 0:
                fenzi +=0
            else:
                fenzi += mu_k[i]*beta_k[i]*(math.log(1+beta_k[i])/math.log(b_k[i]+1))*gu.communication_demand
            fenmu += mu_k[i]*gu.communication_demand
            i+=1
            

        fairness_reward = fenzi/fenmu
        return fairness_reward
        
    def run_episode(self):
        # 模拟一个回合
        for uav in self.uavs:
            # 对于电量进行判断
            # 获取当前状态
            energy_state, cruise_state = self.get_state(uav)

            # # 获取其他无人机的位置    
            other_uavs_positions = [other_uav.position for other_uav in self.uavs if other_uav != uav]
            # 选择动作
            energy_action = uav.select_energy_action(energy_state)
            if uav.energy < 13:
                energy_action = 1
            cruise_action = uav.select_cruise_action(cruise_state,other_uavs_positions,uav.position)

            # 执行动作并获得奖励和下一个状态
            # todo :如果无人机的energy_action选择充电，应该执行寻找最近充电站的操作
            reward_energy,reward, next_energy_state, next_cruise_state = self.perform_action(uav, energy_action, cruise_action)

            
            # 更新Q表
            uav.update_energy_Q(energy_state, energy_action, reward_energy, next_energy_state)
            uav.update_cruise_Q(cruise_state, cruise_action, reward, next_cruise_state)
            print(f"UAV{uav.id} performed action: Energy={energy_action}, Cruise={cruise_action} arrived at the position{uav.position} Cruise_Reward={reward},energy_Reward={reward_energy} remaining energy {uav.energy}")                
